import json
import logging
import random
import time
from paho.mqtt import client as mqtt_client
from typing import Callable, Optional


class MqttClientUtil:
    """MQTT客户端工具类，支持自动订阅和消息发送"""

    def __init__(
            self,
            broker: str = 'broker.emqx.io',
            port: int = 1883,
            topic: str = "python/mqtt",
            client_id: Optional[str] = None,
            username: Optional[str] = None,
            password: Optional[str] = None,
            keepalive: int = 60,
            qos: int = 0,
            reconnect_max_attempts: int = 12,
            reconnect_delay: int = 1
    ):
        """初始化MQTT客户端配置"""
        self.broker = broker
        self.port = port
        self.topic = topic
        self.client_id = client_id or f'python-mqtt-{random.randint(0, 1000)}'
        self.username = username
        self.password = password
        self.keepalive = keepalive
        self.qos = qos
        self.reconnect_max_attempts = reconnect_max_attempts
        self.reconnect_delay = reconnect_delay

        self.client = None
        self._is_connected = False
        self._message_callback = None

        # 配置日志
        self.logger = logging.getLogger("MqttClientUtil")
        self.logger.setLevel(logging.INFO)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        ch = logging.StreamHandler()
        ch.setFormatter(formatter)
        self.logger.addHandler(ch)

    def set_message_callback(self, callback: Callable[[str, dict], None]):
        """设置消息接收回调函数"""
        self._message_callback = callback

    def _on_connect(self, client, userdata, flags, rc):
        """连接成功/失败的回调函数"""
        if rc == 0:
            self._is_connected = True
            self.logger.info(f"Connected to MQTT Broker {self.broker}:{self.port}")
            # 连接成功后订阅主题
            client.subscribe(self.topic, qos=self.qos)
            self.logger.info(f"Subscribed to topic: {self.topic}")
        else:
            self._is_connected = False
            self.logger.error(f"Failed to connect, return code {rc}")

    def _on_disconnect(self, client, userdata, rc):
        """断开连接的回调函数"""
        self._is_connected = False
        self.logger.warning(f"Disconnected from broker, result code: {rc}")

        # 自动重连逻辑
        reconnect_attempts = 0
        while reconnect_attempts < self.reconnect_max_attempts:
            delay = min(self.reconnect_delay * (2 ** reconnect_attempts), 60)
            self.logger.info(f"Reconnecting attempt {reconnect_attempts + 1}/{self.reconnect_max_attempts} in {delay}s")
            time.sleep(delay)

            try:
                client.reconnect()
                self.logger.info("Reconnected successfully")
                return
            except Exception as e:
                reconnect_attempts += 1
                self.logger.error(f"Reconnect attempt {reconnect_attempts} failed: {str(e)}")

        self.logger.critical("Max reconnect attempts reached. MQTT client is offline.")

    def _on_message(self, client, userdata, msg):
        """收到消息的回调函数"""
        try:
            payload = msg.payload.decode()
            self.logger.debug(f"Received message on topic '{msg.topic}': {payload}")

            # 解析JSON消息
            try:
                message_dict = json.loads(payload)
            except json.JSONDecodeError:
                self.logger.warning(f"Failed to parse message as JSON: {payload}")
                message_dict = {"raw": payload}

            # 调用外部回调
            if self._message_callback:
                self._message_callback(msg.topic, message_dict)

        except Exception as e:
            self.logger.error(f"Error processing message: {str(e)}")

    def connect(self):
        """连接到MQTT代理服务器并启动循环"""
        try:
            # 创建MQTT客户端实例
            self.client = mqtt_client.Client(client_id=self.client_id)

            # 设置认证信息（如果有）
            if self.username:
                self.client.username_pw_set(self.username, self.password)

            # 设置回调函数
            self.client.on_connect = self._on_connect
            self.client.on_disconnect = self._on_disconnect
            self.client.on_message = self._on_message

            # 连接并启动循环
            self.client.connect(self.broker, self.port, self.keepalive)
            self.client.loop_start()

            # 等待连接建立
            time.sleep(1)
            if not self._is_connected:
                self.logger.warning("Connection not established within 1 second. Continuing anyway.")

            return True

        except Exception as e:
            self.logger.error(f"Failed to initialize MQTT client: {str(e)}")
            return False

    def disconnect(self):
        """断开与MQTT代理服务器的连接"""
        if self.client and self._is_connected:
            self.client.loop_stop()
            self.client.disconnect()
            self._is_connected = False
            self.logger.info("Disconnected from MQTT broker")

    def publish(self, message: dict, topic: Optional[str] = None):
        """发布消息到MQTT主题"""
        if not self._is_connected:
            self.logger.error("Cannot publish: MQTT client is not connected")
            return False

        try:
            # 使用指定主题或默认主题
            publish_topic = topic or self.topic

            # 转换为JSON字符串
            payload = json.dumps(message)

            # 发布消息
            result = self.client.publish(publish_topic, payload, qos=self.qos)
            status = result[0]

            if status == 0:
                self.logger.info(f"Published message to topic '{publish_topic}': {payload[:100]}")
                return True
            else:
                self.logger.error(f"Failed to publish message to topic '{publish_topic}', status code: {status}")
                return False

        except Exception as e:
            self.logger.error(f"Error publishing message: {str(e)}")
            return False

    def is_connected(self) -> bool:
        """检查MQTT连接状态"""
        return self._is_connected